from AStock.ASSetting import settings, output_path
from AStock.ASQuery import (
    ASQuery_financial,
    ASQuery_nearest_trade_date,
    ASQuery_stock_day,
    ASQuery_stock_info,
    ASQuery_stock_xdxr
)
from AStock.ASUtil import date_to_timestamp, date_to_bar_format
import pandas as pd
from datetime import datetime, timedelta
import os
from pymongo import UpdateOne


class ASGrowthBase(object):
    def __init__(self, report_date, calc_date, df_pool=None):
        now = datetime.now().strftime('%Y%m%d')
        if calc_date is None:
            calc_date = now
        if calc_date == now and datetime.now().hour < 17:
            calc_date = (datetime.now() - timedelta(days=1)).strftime('%Y%m%d')
        begin = int(report_date) - 10000
        self._trade_days = self._get_trade_days(begin, int(now) + 31)
        self._report_date = report_date
        self._calc_date = calc_date
        self._df_pool = df_pool

    def run(self, save_to_db=True, save_to_file=False):
        raise NotImplementedError

    def _get_inputs(self):
        raise NotImplementedError

    def _calc_outputs(self, df_inputs):
        raise NotImplementedError

    def _save(self, df_outputs, to_db=True, to_file=False):
        raise NotImplementedError

    @staticmethod
    def _median(a, b, c):
        x = [a, b, c]
        x = sorted(x)
        return x[1]

    # 取时间较早的
    @staticmethod
    def _early_date(x):
        if pd.isnull(x['financialYjkbDate']):
            v = x['financialYjygDate']
        elif pd.isnull(x['financialYjygDate']):
            v = x['financialYjkbDate']
        elif x['financialYjkbDate'] < x['financialYjygDate']:
            v = x['financialYjkbDate']
        else:
            v = x['financialYjygDate']
        return v if pd.isnull(v) else int(v)

    """
    @staticmethod
    def _get_net_profit(codes, report_date, rename_column):
        df_net_profit = ASQuery_financial(
            code=codes,
            report_date=report_date,
            fields=['code', 'netProfitAfterExtraordinaryGainsAndLosses']
        )
        df_net_profit = df_net_profit.rename(
            columns={'netProfitAfterExtraordinaryGainsAndLosses': rename_column}
        )
        return df_net_profit
    """

    @staticmethod
    def _get_financial_data(codes, report_date, fields):
        """
        查询季报的指定字段数据
        :param codes: list 待查询的股票列表
        :param report_date: str 报告日期
        :param fields: dict 待查询的列英文名: 重命名的列名
        :return: DataFrame 包括列： code 和 重命名的各列
        """
        df_data = ASQuery_financial(
            code=codes,
            report_date=report_date,
            fields=['code'] + list(fields.keys())
        )
        df_data = df_data.rename(
            columns=fields
        )
        return df_data

    @staticmethod
    def _get_trade_days(begin_date, end_date):
        """
        获取交易日
        :param begin_date: int YYYmmdd
        :param end_date: int YYYmmdd
        :return: set of str YYYmmdd
        """
        coll = settings.database.trade_cal
        try:
            cursor = coll.find(
                {
                    'cal_date': {'$gte': begin_date, '$lte': end_date},
                    'is_open': 1
                },
                {
                    '_id': 0,
                    'cal_date': 1
                }
            )
        except Exception as e:
            print(e)
            return None
        data = set([str(c['cal_date']) for c in cursor])
        return data

    def _get_nearest_trade_day(self, date):
        date = datetime.strptime(date, '%Y%m%d')
        max_count = 0
        while max_count < 30:
            str_date = date.strftime('%Y%m%d')
            if str_date in self._trade_days:
                return str_date
            else:
                date -= timedelta(days=1)
            max_count += 1
        return None

    def _fill_total_mv(self, df_notice_date_price):
        df_notice_date_price['date'] = pd.to_datetime(df_notice_date_price['date'])
        if 'total_mv' not in df_notice_date_price.columns:
            df_notice_date_price['total_mv'] = None
        df = df_notice_date_price[pd.isna(df_notice_date_price['total_mv'])]
        df = df.set_index(['date', 'code'], drop=False)
        # df.to_csv('D://df1.csv')
        codes = df['code'].tolist()
        if not codes:
            return df_notice_date_price
        # 使用除权除息数据计算市值
        xdxr = ASQuery_stock_xdxr(codes, fields=['date', 'code', 'shares_after', 'category'])
        # xdxr = xdxr.set_index(['date', 'code'], drop=False)
        mv = xdxr.query('category!=6').loc[:, ['shares_after']].dropna()
        # mv.to_csv('D://mv1.csv')
        df = pd.concat([df, mv], axis=1)
        # df.to_csv('D://df2.csv')
        df = df.assign(
            shares=df.shares_after.groupby(level=1).fillna(method='ffill')
        )
        df['total_mv'] = df.close * df.shares * 10000
        # df.to_csv('D://test11.csv')
        df = df.dropna(subset=['code'])
        df = df[['total_mv']]
        df_notice_date_price = df_notice_date_price.set_index(['date', 'code'], drop=False)
        # df_notice_date_price.to_csv('D://df_notice_date_price1.csv')
        df_notice_date_price = df_notice_date_price.combine_first(df)
        # df_notice_date_price.to_csv('D://df_notice_date_price2.csv')
        return df_notice_date_price

    def _get_price_qfq(self, df_forecast, date_column, offset_days, rename_column, get_total_mv=False):
        """
        获取相对指定日期的股价数据和总市值数据 （前复权）
        :param df_forecast: DataFrame 包含待获取的股票和股价对应日期（通过列名指定）
        :param date_column: str 日期所在的列的列名
        :param offset_days: int 相对预告日的偏移天数，负数-预告日前，正数-预告日后
        :param rename_column: str/tuple  收盘价列命名 总市值列命名
        :param get_total_mv: bool 是否同时获取总市值数据
        :return: DataFrame 相对预告日的日期的股价数据和总市值数据
        """
        projection = {'_id': 0, 'code': 1, 'close': 1, 'date': 1}
        if get_total_mv:
            projection.__setitem__('total_mv', 1)

        or_filter_items = []
        adj_or_filter_items = []
        for tup in df_forecast.itertuples():
            notice_date = tup._asdict()[date_column]
            #print(tup)
            if pd.notna(notice_date):
                notice_date = str(int(notice_date))
                # 计算相对日期
                notice_date = datetime.strptime(notice_date, '%Y%m%d')
                notice_date += timedelta(offset_days)
                notice_date = notice_date.strftime('%Y%m%d')
                # 需要对此查询库，效率低，改为构造时获取过去360天的交易日到本地供查询
                # trade_date = ASQuery_nearest_trade_date(notice_date)
                trade_date = self._get_nearest_trade_day(notice_date)
                # print(notice_date, trade_date)
                # code 和 date_stamp 有索引，检索比 date 字段快
                or_filter_items.append({'code': tup.code, 'date_stamp': date_to_timestamp(trade_date)})
                adj_or_filter_items.append({'code': tup.code, 'date': date_to_bar_format(trade_date)})
        coll = settings.database.stock_day
        try:
            # t1 = datetime.now().timestamp()
            cursor = coll.find({'$or': or_filter_items}, projection)
            # t2 = datetime.now().timestamp()
            # print('or filter num: {} time: {}'.format(len(or_filter_items), t2 - t1))
        except Exception as e:
            print(e)
            return None
        data = [row for row in cursor]

        adj_coll = settings.database.stock_adj
        try:
            adj_cursor = adj_coll.find({'$or': adj_or_filter_items}, {'_id': 0})
        except Exception as e:
            print(e)
            return None
        adj_data = [row for row in adj_cursor]
        df_adj = pd.DataFrame(adj_data, columns=['code', 'date', 'adj'])
        """
        coll = settings.database.stock_day
        data = []
        for tup in df_forecast.itertuples():
            notice_date = str(tup.notice_date)
            # 计算相对日期
            notice_date = datetime.strptime(notice_date, '%Y%m%d')
            notice_date += timedelta(offset_days)
            notice_date = notice_date.strftime('%Y%m%d')
            trade_date = ASQuery_nearest_trade_date(notice_date)
            # code 和 date_stamp 有索引，检索比 date 字段快
            try:
                one = coll.find_one(
                    {'code': tup.code, 'date_stamp': date_to_timestamp(trade_date)},
                    projection
                )
            except Exception as e:
                print(e)
                return None
            if one is None:
                one = {'code': tup.code}
            data.append(one)
        """
        columns = ['code', 'date', 'close']
        if get_total_mv: columns.append('total_mv')
        df = pd.DataFrame(data, columns=columns)
        df = df.join(df_adj.set_index(['code', 'date']), on=['code', 'date'], how='left')
        df['close'] = df['close'] * df['adj']

        if get_total_mv:
            df = self._fill_total_mv(df)

            df = df.rename(columns={'close': rename_column[0], 'total_mv': rename_column[1]})
        else:
            df = df.rename(columns={'close': rename_column})

        df = df.drop(['adj', 'date'], axis=1)
        return df

    def _get_other_quarter_financial_data(self, codes, df, income=True, profit=True):
        # 去年同期扣非净利润
        fields = {}
        if profit:
            fields.__setitem__('netProfitAfterExtraordinaryGainsAndLosses', 'lastYearSameQuarterNonNetProfit')
        if income:
            fields.__setitem__('totalOperatingRevenue', 'lastYearSameQuarterOperatingIncome')
        last_year_same_quarter_report_date = str(int(self._report_date) - 10000)
        df_last_year_same_quarter_net_profit = self._get_financial_data(
            codes,
            last_year_same_quarter_report_date,
            fields
        )
        df = df.join(df_last_year_same_quarter_net_profit.set_index('code'), on='code', how='left')

        # 今年上一季度报告扣非净利润
        if self._report_date[4:6] == '03':
            df['thisYearLastQuarterNonNetProfit'] = 0
            df['thisYearLastQuarterOperatingIncome'] = 0
        else:
            if self._report_date[4:6] == '06':
                this_year_last_quarter_report_date = self._report_date[:4] + '0331'
            elif self._report_date[4:6] == '09':
                this_year_last_quarter_report_date = self._report_date[:4] + '0630'
            elif self._report_date[4:6] == '12':
                this_year_last_quarter_report_date = self._report_date[:4] + '0930'
            else:
                assert 0

            fields = {}
            if profit:
                fields.__setitem__('netProfitAfterExtraordinaryGainsAndLosses', 'thisYearLastQuarterNonNetProfit')
            if income:
                fields.__setitem__('totalOperatingRevenue', 'thisYearLastQuarterOperatingIncome')

            df_this_year_last_quarter_net_profit = self._get_financial_data(
                codes,
                this_year_last_quarter_report_date,
                fields
            )
            df = df.join(df_this_year_last_quarter_net_profit.set_index('code'), on='code', how='left')

        # 去年上一季度报告扣非净利润
        last_year_report_date = str(int(self._report_date) - 10000)
        if last_year_report_date[4:6] == '03':
            df['lastYearLastQuarterNonNetProfit'] = 0
            df['lastYearLastQuarterOperatingIncome'] = 0
        else:
            if last_year_report_date[4:6] == '06':
                last_year_last_quarter_report_date = last_year_report_date[:4] + '0331'
            elif last_year_report_date[4:6] == '09':
                last_year_last_quarter_report_date = last_year_report_date[:4] + '0630'
            elif last_year_report_date[4:6] == '12':
                last_year_last_quarter_report_date = last_year_report_date[:4] + '0930'
            else:
                assert 0

            fields = {}
            if profit:
                fields.__setitem__('netProfitAfterExtraordinaryGainsAndLosses', 'lastYearLastQuarterNonNetProfit')
            if income:
                fields.__setitem__('totalOperatingRevenue', 'lastYearLastQuarterOperatingIncome')

            df_last_year_last_quarter_net_profit = self._get_financial_data(
                codes,
                last_year_last_quarter_report_date,
                fields
            )
            df = df.join(df_last_year_last_quarter_net_profit.set_index('code'), on='code', how='left')

        # 今年上上一季度报告扣非净利润
        if self._report_date[4:6] in ['03', '06']:
            df['thisYearLastLastQuarterNonNetProfit'] = 0
            df['thisYearLastLastQuarterOperatingIncome'] = 0
        else:
            if self._report_date[4:6] == '09':
                this_year_last_last_quarter_report_date = self._report_date[:4] + '0331'
            elif self._report_date[4:6] == '12':
                this_year_last_last_quarter_report_date = self._report_date[:4] + '0630'
            else:
                assert 0

            fields = {}
            if profit:
                fields.__setitem__('netProfitAfterExtraordinaryGainsAndLosses', 'thisYearLastLastQuarterNonNetProfit')
            if income:
                fields.__setitem__('totalOperatingRevenue', 'thisYearLastLastQuarterOperatingIncome')

            df_this_year_last_last_quarter_net_profit = self._get_financial_data(
                codes,
                this_year_last_last_quarter_report_date,
                fields
            )
            df = df.join(df_this_year_last_last_quarter_net_profit.set_index('code'), on='code', how='left')

        # 去年年报扣非净利润
        fields = {}
        if profit:
            fields.__setitem__('netProfitAfterExtraordinaryGainsAndLosses', 'lastYearAnnualNonNetProfit')
        if income:
            fields.__setitem__('totalOperatingRevenue', 'lastYearAnnualOperatingIncome')
        last_year_annual_report_date = str(int(self._report_date[:4]) - 1) + '1231'
        df_last_year_annual_net_profit = self._get_financial_data(
            codes,
            last_year_annual_report_date,
            fields
        )
        df = df.join(df_last_year_annual_net_profit.set_index('code'), on='code', how='left')

        # 去年3季报扣非净利润
        fields = {}
        if profit:
            fields.__setitem__('netProfitAfterExtraordinaryGainsAndLosses', 'lastYearQuarter3NonNetProfit')
        if income:
            fields.__setitem__('totalOperatingRevenue', 'lastYearQuarter3OperatingIncome')
        last_year_q3_report_date = str(int(self._report_date[:4]) - 1) + '0930'
        df_last_year_q3_net_profit = self._get_financial_data(
            codes,
            last_year_q3_report_date,
            fields
        )
        df = df.join(df_last_year_q3_net_profit.set_index('code'), on='code', how='left')

        return df


class ASSingleGrowth(ASGrowthBase):
    def __init__(self, report_date, calc_date=None, df_pool=None):
        super(ASSingleGrowth, self).__init__(report_date, calc_date=calc_date, df_pool=df_pool)

    def run(self, save_to_db=True, save_to_file=False):
        print('single growth {} running...'.format(self._report_date))
        df_inputs = self._get_inputs()
        if df_inputs is None:
            print('error calculating single growth: failed to get input data')
            return
        df_outputs = self._calc_outputs(df_inputs)
        if df_outputs is None:
            print('error calculating single growth: failed to calculate output data')
            return
        self._save(df_outputs, to_db=save_to_db, to_file=save_to_file)

    def _get_inputs(self):
        """
        获取预报盈利数据
        :return: DataFrame, columns: [name, industry,
                        code, report_date, notice_date, forecastNetProfit,
                        noticeDatePrice, calcDateTotalMV, monthBeforeNoticeDatePrice,
                        weekAfterNoticeDatePrice, calcDatePrice, lastYearSameQuarterNonNetProfit,
                        thisYearLastQuarterNonNetProfit, lastYearLastQuarterNonNetProfit,
                        thisYearLastLastQuarterNonNetProfit, lastYearAnnualNonNetProfit,
                        lastYearQuarter3NonNetProfit, changeReason]
                lastYearQuarter3NonNetProfit - 去年3季报，用于今年一季计算环比增长率
        """
        fields = [
            'code',
            'report_date',
            'financialYjkbDate',
            'financialYjygDate',
            'forecastNetProfit',
            'changeReason'
        ]
        # 如果传入了股票池股票列表，则只获取股票池的数据
        if self._df_pool is not None:
            codes = self._df_pool['code'].tolist()
        else:
            codes = None

        df = ASQuery_financial(code=codes, report_date=self._report_date, fields=fields)
        # 排除 NaN 和 null
        df = df.dropna(subset=['forecastNetProfit'])
        df = df.dropna(how='all', subset=['financialYjygDate', 'financialYjkbDate'])
        if df.empty:
            return None

        # df = df.head(5)

        df_info = ASQuery_stock_info(code=df['code'].tolist(), fields=['code', 'name', 'industry'])
        df = df.join(df_info.set_index('code'), on='code', how='left')

        df['notice_date'] = df.apply(lambda x: self._early_date(x), axis=1)
        df = df.drop('financialYjkbDate', axis=1)
        df = df.drop('financialYjygDate', axis=1)

        codes = df['code'].tolist()

        # 获取预告日（前一日）股价数据
        # TODO 此处想要获取的是发布报告前的股价，如何考虑报告盘前发布影响？
        # 一般都是预告日的前一日盘后发布（即预告日盘前），所以取预告日前一天收盘价
        df_notice_date_price = self._get_price_qfq(df, 'notice_date', -1, 'noticeDatePrice')
        df = df.join(df_notice_date_price.set_index('code'), on='code', how='left')

        # 预告日期前一月收盘价
        df_month_before_price = self._get_price_qfq(df, 'notice_date', -30, 'monthBeforeNoticeDatePrice')
        df = df.join(df_month_before_price.set_index('code'), on='code', how='left')

        # 预告日期一周后收盘价
        df_week_after_price = self._get_price_qfq(df, 'notice_date', 7, 'weekAfterNoticeDatePrice')
        df = df.join(df_week_after_price.set_index('code'), on='code', how='left')

        # 计算日的收盘价和市值数据
        date = ASQuery_nearest_trade_date(self._calc_date)
        df_now_price = ASQuery_stock_day(codes, date, date, fields=['code', 'close', 'total_mv'], type='qfq')
        df_now_price = df_now_price.rename(columns={'close': 'calcDatePrice', 'total_mv': 'calcDateTotalMV'})
        df = df.join(df_now_price.set_index('code'), on='code', how='left')

        # 其他季度的扣非净利数据
        df = self._get_other_quarter_financial_data(codes, df, income=False, profit=True)

        return df

    def _calc_outputs(self, df_inputs):
        """
        计算单增输出数据
        :param df_inputs: input DataFrame, columns: [name, industry,
                        code, report_date, notice_date, forecastNetProfit,
                        noticeDatePrice, calcDateTotalMV, monthBeforeNoticeDatePrice,
                        weekAfterNoticeDatePrice, calcDatePrice, lastYearSameQuarterNonNetProfit,
                        thisYearLastQuarterNonNetProfit, lastYearLastQuarterNonNetProfit,
                        thisYearLastLastQuarterNonNetProfit, lastYearAnnualNonNetProfit,
                        lastYearQuarter3NonNetProfit, changeReason]
        :return: DataFrame, columns: df_inputs columns +
                        [yoyTotalGrowthRate, yoyQuarterGrowthRate,
                        qoqGrowthRate, deductPe, deductPeg, percentChangeToCalcDate,
                        percentChangeAfterWeek, percentChangeBeforeMonth]
        """
        # 累计同比增速=(预告扣非净利润-去年同期扣非净利润)/去年同期扣非净利润
        df_outputs = df_inputs.assign(
            yoyTotalGrowthRate=(df_inputs['forecastNetProfit'] - df_inputs['lastYearSameQuarterNonNetProfit']) /
                               df_inputs[
                                   'lastYearSameQuarterNonNetProfit'].abs() * 100)
        # 单季同比增速=[(预告扣非净利润-今年上一季度报告扣非净利润)-(去年同期扣非净利润-去年上一季度报告扣非净利润)]
        #               /(去年同期扣非净利润-去年上一季度报告扣非净利润)
        df_outputs = df_outputs.assign(
            yoyQuarterGrowthRate=((df_outputs['forecastNetProfit'] - df_outputs['thisYearLastQuarterNonNetProfit'])
                                  - (df_outputs['lastYearSameQuarterNonNetProfit'] - df_outputs[
                        'lastYearLastQuarterNonNetProfit']))
                                 / (df_outputs['lastYearSameQuarterNonNetProfit'] - df_outputs[
                'lastYearLastQuarterNonNetProfit']).abs() * 100)
        # 单季环比增速=[(预告扣非净利润-今年上一季度报告扣非净利润)-(今年上一季度报告扣非净利润-今年上上一季度报告扣非净利润)]
        #               /(今年上一季度报告扣非净利润-今年上上一季度报告扣非净利润)
        if self._report_date[4:6] == '03':
            df_outputs = df_outputs.assign(
                qoqGrowthRate=(df_outputs['forecastNetProfit']
                               - (df_outputs['lastYearAnnualNonNetProfit'] - df_outputs[
                            'lastYearQuarter3NonNetProfit']))
                              / (df_outputs['lastYearAnnualNonNetProfit'] - df_outputs[
                    'lastYearQuarter3NonNetProfit']).abs() * 100)
        else:
            df_outputs = df_outputs.assign(
                qoqGrowthRate=((df_outputs['forecastNetProfit'] - df_outputs['thisYearLastQuarterNonNetProfit'])
                               - (df_outputs['thisYearLastQuarterNonNetProfit'] - df_outputs[
                            'thisYearLastLastQuarterNonNetProfit']))
                              / (df_outputs['thisYearLastQuarterNonNetProfit'] - df_outputs[
                    'thisYearLastLastQuarterNonNetProfit']).abs() * 100)
        # 扣非PE=预告日期总市值/(预告扣非净利润+去年年报扣非净利润-去年同期扣非净利润)
        df_outputs = df_outputs.assign(
            deductPe=pd.to_numeric(df_outputs['calcDateTotalMV'] / (df_outputs[
                                                            'forecastNetProfit'] + df_outputs[
                                                            'lastYearAnnualNonNetProfit'] - df_outputs[
                                                            'lastYearSameQuarterNonNetProfit'])))

        # 扣非PEG=扣非PE/100*MEDIAN(年度累计同比增速, 单季同比增速, 单季环比增速)
        # pe和g都为负数，则peg应该为负数
        df_outputs['deductPeg'] = df_outputs.apply(
            lambda x: x['deductPe'] / (
                self._median(x['yoyQuarterGrowthRate'], x['qoqGrowthRate'], x['yoyTotalGrowthRate'])) * (
                          -1 if x['deductPe'] < 0 and self._median(x['yoyQuarterGrowthRate'], x['qoqGrowthRate'],
                                                                   x['yoyTotalGrowthRate']) < 0 else 1), axis=1)

        # 至今涨跌幅
        df_outputs = df_outputs.assign(
            percentChangeToCalcDate=100 * (df_outputs['calcDatePrice'] - df_outputs['noticeDatePrice']) / df_outputs[
                'noticeDatePrice'])
        # 一周后涨跌幅
        df_outputs = df_outputs.assign(
            percentChangeAfterWeek=100 * (df_outputs['weekAfterNoticeDatePrice'] - df_outputs['noticeDatePrice']) /
                                   df_outputs['noticeDatePrice'])
        # 披露前一月涨跌幅
        df_outputs = df_outputs.assign(
            percentChangeBeforeMonth=100 * (df_outputs['noticeDatePrice'] - df_outputs['monthBeforeNoticeDatePrice']) /
                                     df_outputs['monthBeforeNoticeDatePrice'])
        df_outputs = df_outputs.round(2)
        return df_outputs

    def _save(self, df_outputs, to_db=True, to_file=False):
        columns = ['code', 'name', 'industry', 'report_date', 'notice_date',
                   'forecastNetProfit', 'noticeDatePrice',
                   'calcDateTotalMV', 'monthBeforeNoticeDatePrice',
                   'weekAfterNoticeDatePrice', 'calcDatePrice',
                   'lastYearSameQuarterNonNetProfit', 'thisYearLastQuarterNonNetProfit',
                   'lastYearLastQuarterNonNetProfit', 'thisYearLastLastQuarterNonNetProfit',
                   'lastYearAnnualNonNetProfit', 'lastYearQuarter3NonNetProfit',
                   'yoyTotalGrowthRate', 'yoyQuarterGrowthRate',
                   'qoqGrowthRate', 'deductPe', 'deductPeg', 'percentChangeToCalcDate',
                   'percentChangeAfterWeek', 'percentChangeBeforeMonth', 'changeReason']
        if self._df_pool is not None:
            columns.append('groups')
            df_outputs = df_outputs.join(self._df_pool.set_index('code'), on='code', how='left')

        if to_file:
            if self._df_pool is not None:
                out_put_file = os.path.join(output_path, 'pool_single_growth_{}_{}.xlsx'.format(
                    self._report_date, datetime.now().strftime('%Y%m%d')))
            else:
                out_put_file = os.path.join(output_path, 'single_growth_{}_{}.xlsx'.format(
                    self._report_date, datetime.now().strftime('%Y%m%d')))

            renamed_columns = {
                'name': '名称', 'code': '代码', 'industry': '行业',
                'notice_date': '公告日期',
                'calcDatePrice': '当前价格', 'calcDateTotalMV': '当前市值',
                'forecastNetProfit': '预测净利', 'lastYearSameQuarterNonNetProfit': '去年同期扣非净利',
                'lastYearAnnualNonNetProfit': '去年年报扣非净利',
                'yoyTotalGrowthRate': '扣非净利累计同比',
                'yoyQuarterGrowthRate': '扣非净利单季同比',
                'qoqGrowthRate': '扣非净利季度环比',
                'deductPe': '扣非PE', 'deductPeg': '扣非PEG',
                'percentChangeToCalcDate': '公告至今涨幅',
                'percentChangeAfterWeek': '公告一周后涨幅',
                'percentChangeBeforeMonth': '公告一月前涨幅',
                'changeReason': '变动原因'
            }
            if 'groups' in df_outputs.columns:
                renamed_columns.__setitem__('groups', '分组')
            df_outputs = df_outputs.rename(columns=renamed_columns)

            df_outputs.to_excel(
                out_put_file,
                index=False,
                encoding='utf-8',
                columns=[renamed_columns[c] for c in columns if c in renamed_columns]
            )
            print('saved single growth to {}'.format(out_put_file))

        if to_db:
            bulk = []
            for tup in df_outputs.itertuples(index=False):
                one = UpdateOne(
                    {'code': tup.code, 'report_date': tup.report_date},
                    {'$set': tup._asdict()},
                    upsert=True
                )
                bulk.append(one)
            try:
                coll = settings.database.single_growth
                coll.bulk_write(bulk)
            except Exception as e:
                print(e)
            print('saved {} items to single_growth'.format(len(bulk)))


class ASDoubleGrowth(ASGrowthBase):
    def __init__(self, report_date, calc_date=None, df_pool=None):
        super(ASDoubleGrowth, self).__init__(report_date, calc_date=calc_date, df_pool=df_pool)

    def run(self, save_to_db=True, save_to_file=False):
        print('double growth {} running...'.format(self._report_date))
        df_inputs = self._get_inputs()
        if df_inputs is None:
            print('error calculating double growth: failed to get input data')
            return
        df_outputs = self._calc_outputs(df_inputs)
        if df_outputs is None:
            print('error calculating double growth: failed to calculate output data')
            return
        self._save(df_outputs, to_db=save_to_db, to_file=save_to_file)

    def _get_inputs(self):
        """
        获取扣非盈利数据和营业总收入数据
        :return: DataFrame 主要包含“日期、价格数据”、“扣非净利数据”、“营业收入数据”
                columns: [name, industry, code, report_date, calcDatePrice,
                        forecast_date, forecastDatePrice, monthBeforeForecastDatePrice, weekAfterForecastDatePrice,
                        notice_date, noticeDatePrice, calcDateTotalMV, monthBeforeNoticeDatePrice, weekAfterNoticeDatePrice,
                        nonNetProfit, thisYearLastLastQuarterNonNetProfit, thisYearLastQuarterNonNetProfit,
                        lastYearSameQuarterNonNetProfit, lastYearLastQuarterNonNetProfit, lastYearAnnualNonNetProfit, lastYearQuarter3NonNetProfit,
                        operatingIncome, thisYearLastLastQuarterOperatingIncome, thisYearLastQuarterOperatingIncome,
                        lastYearSameQuarterOperatingIncome, lastYearLastQuarterOperatingIncome, lastYearAnnualOperatingIncome, lastYearQuarter3OperatingIncome]
                lastYearQuarter3NonNetProfit - 去年3季报扣非净利，用于今年一季计算环比增长率
                lastYearQuarter3OperatingIncome - 去年3季报营收，用于今年一季计算环比增长率
        """
        # 获取财务报表中本季度的数据
        fields = [
            'code',
            'report_date',
            'financialYjkbDate',
            'financialYjygDate',
            'financialNoticeDate',
            'totalOperatingRevenue',  # operatingIncome
            'netProfitAfterExtraordinaryGainsAndLosses'  # nonNetProfit
        ]
        # 如果传入了股票池股票列表，则只获取股票池的数据
        if self._df_pool is not None:
            codes = self._df_pool['code'].tolist()
        else:
            codes = None

        df = ASQuery_financial(code=codes, report_date=self._report_date, fields=fields)
        # 排除 NaN 和 null
        row_count = df.shape[0]
        df = df.dropna(
            how='any',
            subset=['financialNoticeDate', 'totalOperatingRevenue', 'netProfitAfterExtraordinaryGainsAndLosses']
        )

        df = df[~df['financialNoticeDate'].isin([0])]

        new_row_count = df.shape[0]
        if new_row_count < row_count:
            print("========= {} stock's financial data missing =========".format(row_count - new_row_count))

        if df.empty:
            return None

        # df = df.head(5)

        # 获取行业等基础数据
        df_info = ASQuery_stock_info(code=df['code'].tolist(), fields=['code', 'name', 'industry'])
        df = df.join(df_info.set_index('code'), on='code', how='left')

        df['forecast_date'] = df.apply(lambda x: self._early_date(x), axis=1)
        df = df.drop('financialYjkbDate', axis=1)
        df = df.drop('financialYjygDate', axis=1)

        df = df.rename(columns={
            'financialNoticeDate': 'notice_date',
            'totalOperatingRevenue': 'operatingIncome',
            'netProfitAfterExtraordinaryGainsAndLosses': 'nonNetProfit'})

        codes = df['code'].tolist()

        df = self._get_all_price_data(codes, df)

        # 其他季度的扣非净利和营业总收入数据
        df = self._get_other_quarter_financial_data(codes, df, income=True, profit=True)

        return df

    def _get_all_price_data(self, codes, df):
        # 获取预告日（前一日）股价数据
        # TODO 此处想要获取的是发布报告前的股价，如何考虑报告盘前发布影响？
        # 一般都是预告日的前一日盘后发布（即预告日盘前），所以取预告日前一天收盘价
        df_forecast_date_price = self._get_price_qfq(
            df, 'forecast_date', -1, 'forecastDatePrice')
        df = df.join(df_forecast_date_price.set_index('code'), on='code', how='left')

        # 预告日期前一月收盘价
        df_month_before_forecast_price = self._get_price_qfq(df, 'forecast_date', -30, 'monthBeforeForecastDatePrice')
        df = df.join(df_month_before_forecast_price.set_index('code'), on='code', how='left')

        # 预告日期一周后收盘价
        df_week_after_forecast_price = self._get_price_qfq(df, 'forecast_date', 7, 'weekAfterForecastDatePrice')
        df = df.join(df_week_after_forecast_price.set_index('code'), on='code', how='left')

        # 获取财报公告日股价数据
        # TODO 此处想要获取的是发布报告前的股价，如何考虑报告盘前发布影响？
        # 一般都是公告日的前一日盘后发布（即公告日盘前），所以取公告日前一天收盘价
        df_notice_date_price = self._get_price_qfq(df, 'notice_date', -1, 'noticeDatePrice')
        # print(df_notice_date_price.columns)
        # print(df_notice_date_price.head(10))
        df = df.join(df_notice_date_price.set_index('code'), on='code', how='left')

        # 财报公告日期前一月收盘价
        df_month_before_price = self._get_price_qfq(df, 'notice_date', -30, 'monthBeforeNoticeDatePrice')
        df = df.join(df_month_before_price.set_index('code'), on='code', how='left')

        # 财报公告日期一周后收盘价
        df_week_after_price = self._get_price_qfq(df, 'notice_date', 7, 'weekAfterNoticeDatePrice')
        df = df.join(df_week_after_price.set_index('code'), on='code', how='left')

        # 计算日的收盘价和市值
        date = ASQuery_nearest_trade_date(self._calc_date)
        df_now_price = ASQuery_stock_day(codes, date, date, fields=['code', 'close', 'total_mv'], type='qfq')
        df_now_price = df_now_price.rename(columns={'close': 'calcDatePrice', 'total_mv': 'calcDateTotalMV'})
        df = df.join(df_now_price.set_index('code'), on='code', how='left')

        return df

    def _calc_outputs(self, df_inputs):
        """
        计算双增输出数据
        :param df_inputs: DataFrame 主要包含“日期、价格数据”、“扣非净利数据”、“营业收入数据”
                    columns: [name, industry, code, report_date, calcDatePrice,
                            forecast_date, forecastDatePrice, monthBeforeForecastDatePrice, weekAfterForecastDatePrice,
                            notice_date, noticeDatePrice, calcDateTotalMV, monthBeforeNoticeDatePrice, weekAfterNoticeDatePrice,
                            nonNetProfit, thisYearLastLastQuarterNonNetProfit, thisYearLastQuarterNonNetProfit,
                            lastYearSameQuarterNonNetProfit, lastYearLastQuarterNonNetProfit, lastYearAnnualNonNetProfit, lastYearQuarter3NonNetProfit,
                            operatingIncome, thisYearLastLastQuarterOperatingIncome, thisYearLastQuarterOperatingIncome,
                            lastYearSameQuarterOperatingIncome, lastYearLastQuarterOperatingIncome, lastYearAnnualOperatingIncome, lastYearQuarter3OperatingIncome]
                    lastYearQuarter3NonNetProfit - 去年3季报扣非净利，用于今年一季计算环比增长率
                    lastYearQuarter3OperatingIncome - 去年3季报营收，用于今年一季计算环比增长率
        :return: DataFrame, columns: df_inputs columns +
                        ['yoyProfitTotalGrowthRate', 'yoyProfitQuarterGrowthRate',
                       'qoqProfitGrowthRate', 'deductPe', 'deductPeg',
                       'yoyIncomeTotalGrowthRate', 'yoyIncomeQuarterGrowthRate',
                       'qoqIncomeGrowthRate', 'deductPegIncome',
                       'percentChangeToCalcDate', 'percentChangeAfterWeek', 'percentChangeBeforeMonth']
        """
        # 累计同比增速=(预告扣非净利润-去年同期扣非净利润)/去年同期扣非净利润
        df_outputs = df_inputs.assign(
            yoyProfitTotalGrowthRate=(df_inputs['nonNetProfit'] - df_inputs['lastYearSameQuarterNonNetProfit']) /
                                     df_inputs[
                                         'lastYearSameQuarterNonNetProfit'].abs() * 100)
        df_outputs = df_outputs.assign(
            yoyIncomeTotalGrowthRate=(df_inputs['operatingIncome'] - df_inputs['lastYearSameQuarterOperatingIncome']) /
                                     df_inputs[
                                         'lastYearSameQuarterOperatingIncome'].abs() * 100)
        # 单季同比增速=[(预告扣非净利润-今年上一季度报告扣非净利润)-(去年同期扣非净利润-去年上一季度报告扣非净利润)]
        #               /(去年同期扣非净利润-去年上一季度报告扣非净利润)
        df_outputs = df_outputs.assign(
            yoyProfitQuarterGrowthRate=((df_outputs['nonNetProfit'] - df_outputs['thisYearLastQuarterNonNetProfit'])
                                        - (df_outputs['lastYearSameQuarterNonNetProfit'] - df_outputs[
                        'lastYearLastQuarterNonNetProfit']))
                                       / (df_outputs['lastYearSameQuarterNonNetProfit'] - df_outputs[
                'lastYearLastQuarterNonNetProfit']).abs() * 100)
        df_outputs = df_outputs.assign(
            yoyIncomeQuarterGrowthRate=((df_outputs['operatingIncome'] - df_outputs[
                'thisYearLastQuarterOperatingIncome'])
                                        - (df_outputs['lastYearSameQuarterOperatingIncome'] - df_outputs[
                        'lastYearLastQuarterOperatingIncome']))
                                       / (df_outputs['lastYearSameQuarterOperatingIncome'] - df_outputs[
                'lastYearLastQuarterOperatingIncome']).abs() * 100)
        # 单季环比增速=[(预告扣非净利润-今年上一季度报告扣非净利润)-(今年上一季度报告扣非净利润-今年上上一季度报告扣非净利润)]
        #               /(今年上一季度报告扣非净利润-今年上上一季度报告扣非净利润)
        if self._report_date[4:6] == '03':
            df_outputs = df_outputs.assign(
                qoqProfitGrowthRate=(df_outputs['nonNetProfit']
                                     - (df_outputs['lastYearAnnualNonNetProfit'] - df_outputs[
                            'lastYearQuarter3NonNetProfit']))
                                    / (df_outputs['lastYearAnnualNonNetProfit'] - df_outputs[
                    'lastYearQuarter3NonNetProfit']).abs() * 100)
            df_outputs = df_outputs.assign(
                qoqIncomeGrowthRate=(df_outputs['operatingIncome']
                                     - (df_outputs['lastYearAnnualOperatingIncome'] - df_outputs[
                            'lastYearQuarter3OperatingIncome']))
                                    / (df_outputs['lastYearAnnualOperatingIncome'] - df_outputs[
                    'lastYearQuarter3OperatingIncome']).abs() * 100)
        else:
            df_outputs = df_outputs.assign(
                qoqProfitGrowthRate=((df_outputs['nonNetProfit'] - df_outputs['thisYearLastQuarterNonNetProfit'])
                                     - (df_outputs['thisYearLastQuarterNonNetProfit'] - df_outputs[
                            'thisYearLastLastQuarterNonNetProfit']))
                                    / (df_outputs['thisYearLastQuarterNonNetProfit'] - df_outputs[
                    'thisYearLastLastQuarterNonNetProfit']).abs() * 100)
            df_outputs = df_outputs.assign(
                qoqIncomeGrowthRate=((df_outputs['operatingIncome'] - df_outputs['thisYearLastQuarterOperatingIncome'])
                                     - (df_outputs['thisYearLastQuarterOperatingIncome'] - df_outputs[
                            'thisYearLastLastQuarterOperatingIncome']))
                                    / (df_outputs['thisYearLastQuarterOperatingIncome'] - df_outputs[
                    'thisYearLastLastQuarterOperatingIncome']).abs() * 100)
        # 扣非PE=预告日期总市值/(预告扣非净利润+去年年报扣非净利润-去年同期扣非净利润)
        df_outputs = df_outputs.assign(
            deductPe=pd.to_numeric(df_outputs['calcDateTotalMV'] / (df_outputs[
                                                            'nonNetProfit'] + df_outputs[
                                                            'lastYearAnnualNonNetProfit'] - df_outputs[
                                                            'lastYearSameQuarterNonNetProfit'])))
        df_outputs = df_outputs.assign(
            deductPeIncome=pd.to_numeric(df_outputs['calcDateTotalMV'] / (df_outputs[
                                                            'operatingIncome'] + df_outputs[
                                                            'lastYearAnnualOperatingIncome'] - df_outputs[
                                                            'lastYearSameQuarterOperatingIncome'])))


        # 扣非PEG=扣非PE/100*MEDIAN(年度累计同比增速, 单季同比增速, 单季环比增速)
        # pe和g都为负数，则peg应该为负数
        """
        df_outputs['deductPeg'] = df_outputs.apply(
            lambda x: x['deductPe'] / (
                self._median(x['yoyProfitQuarterGrowthRate'], x['qoqProfitGrowthRate'], x['yoyProfitTotalGrowthRate'])) * (
                          -1 if x['deductPe'] < 0 and self._median(x['yoyProfitQuarterGrowthRate'], x['qoqProfitGrowthRate'],
                                                                   x['yoyProfitTotalGrowthRate']) < 0 else 1), axis=1)
        df_outputs['deductPegIncome'] = df_outputs.apply(
            lambda x: x['deductPe'] / (
                self._median(x['yoyIncomeQuarterGrowthRate'], x['qoqIncomeGrowthRate'],
                             x['yoyIncomeTotalGrowthRate'])) * (
                          -1 if x['deductPe'] < 0 and self._median(x['yoyIncomeQuarterGrowthRate'],
                                                                   x['qoqIncomeGrowthRate'],
                                                                   x['yoyIncomeTotalGrowthRate']) < 0 else 1), axis=1)
        """
        # 使用单季度增速计算即期peg
        df_outputs['deductPeg'] = df_outputs.apply(
            lambda x: x['deductPe'] / (
                x['yoyProfitQuarterGrowthRate']) * (
                          -1 if x['deductPe'] < 0 and x['yoyProfitQuarterGrowthRate'] < 0 else 1), axis=1)
        df_outputs['deductPegIncome'] = df_outputs.apply(
            lambda x: x['deductPe'] / (
                x['yoyIncomeQuarterGrowthRate']) * (
                          -1 if x['deductPe'] < 0 and x['yoyIncomeQuarterGrowthRate'] < 0 else 1), axis=1)

        # 选择预报或财报公告日期较早的日期收盘价作为参考价格
        df_outputs['firstNoticeDataPrice'] = df_outputs['forecastDatePrice'].combine_first(
            df_outputs['noticeDatePrice'])
        df_outputs['weekAfterFirstNoticeDataPrice'] = df_outputs['weekAfterForecastDatePrice'].combine_first(
            df_outputs['weekAfterNoticeDatePrice'])
        df_outputs['monthBeforeFirstNoticeDataPrice'] = df_outputs['monthBeforeForecastDatePrice'].combine_first(
            df_outputs['monthBeforeNoticeDatePrice'])

        # （预报、财报最早披露日期）至今涨跌幅
        df_outputs = df_outputs.assign(
            percentChangeToCalcDate=100 * (df_outputs['calcDatePrice'] - df_outputs['firstNoticeDataPrice']) / df_outputs[
                'firstNoticeDataPrice'])
        # （预报、财报最早披露日期）一周后涨跌幅
        df_outputs = df_outputs.assign(
            percentChangeAfterWeek=100 * (
                        df_outputs['weekAfterFirstNoticeDataPrice'] - df_outputs['firstNoticeDataPrice']) /
                                   df_outputs['firstNoticeDataPrice'])
        # （预报、财报最早披露日期）披露前一月涨跌幅
        df_outputs = df_outputs.assign(
            percentChangeBeforeMonth=100 * (
                        df_outputs['firstNoticeDataPrice'] - df_outputs['monthBeforeFirstNoticeDataPrice']) /
                                     df_outputs['monthBeforeFirstNoticeDataPrice'])

        df_outputs = df_outputs.drop(
            ['firstNoticeDataPrice', 'weekAfterFirstNoticeDataPrice', 'monthBeforeFirstNoticeDataPrice'], axis=1)
        df_outputs = df_outputs.round(2)
        return df_outputs

    def _save(self, df_outputs, to_db=True, to_file=False):
        columns = ['code', 'name', 'industry', 'forecast_date', 'notice_date',

                   'yoyProfitTotalGrowthRate', 'yoyProfitQuarterGrowthRate',
                   'qoqProfitGrowthRate', 'deductPe', 'deductPeg',

                   'yoyIncomeTotalGrowthRate', 'yoyIncomeQuarterGrowthRate',
                   'qoqIncomeGrowthRate', 'deductPegIncome',

                   'percentChangeToCalcDate', 'percentChangeAfterWeek', 'percentChangeBeforeMonth',

                   'report_date', 'forecastDatePrice', 'noticeDatePrice', 'calcDatePrice', 'calcDateTotalMV',
                   'monthBeforeForecastDatePrice', 'weekAfterForecastDatePrice',
                   'monthBeforeNoticeDatePrice', 'weekAfterNoticeDatePrice',

                   'nonNetProfit', 'thisYearLastLastQuarterNonNetProfit', 'thisYearLastQuarterNonNetProfit',
                   'lastYearSameQuarterNonNetProfit', 'lastYearLastQuarterNonNetProfit', 'lastYearAnnualNonNetProfit',
                   'lastYearQuarter3NonNetProfit',
                   'operatingIncome', 'thisYearLastLastQuarterOperatingIncome', 'thisYearLastQuarterOperatingIncome',
                   'lastYearSameQuarterOperatingIncome', 'lastYearLastQuarterOperatingIncome',
                   'lastYearAnnualOperatingIncome', 'lastYearQuarter3OperatingIncome']

        if self._df_pool is not None:
            columns.append('groups')
            df_outputs = df_outputs.join(self._df_pool.set_index('code'), on='code', how='left')

        if to_file:
            if self._df_pool is not None:
                out_put_file = os.path.join(output_path, 'pool_double_growth_{}_{}.xlsx'.format(
                    self._report_date, datetime.now().strftime('%Y%m%d')))
            else:
                out_put_file = os.path.join(output_path, 'double_growth_{}_{}.xlsx'.format(
                    self._report_date, datetime.now().strftime('%Y%m%d')))
            renamed_columns = {
                'name': '名称', 'code': '代码', 'industry': '行业',
                'forecast_date': '快报日期', 'notice_date': '公告日期',
                'calcDatePrice': '当前价格', 'calcDateTotalMV': '当前市值',
                'nonNetProfit': '扣非净利', 'lastYearSameQuarterNonNetProfit': '去年同期扣非净利',
                'lastYearAnnualNonNetProfit': '去年年报扣非净利', 'lastYearLastQuarterNonNetProfit': '去年上季扣非净利',
                'operatingIncome': '营收', 'lastYearSameQuarterOperatingIncome': '去年同期营收',
                'lastYearAnnualOperatingIncome': '去年年报营收', 'lastYearLastQuarterOperatingIncome': '去年上季营收',
                'yoyProfitTotalGrowthRate': '扣非净利累计同比',
                'yoyProfitQuarterGrowthRate': '扣非净利单季同比',
                'qoqProfitGrowthRate': '扣非净利季度环比',
                'deductPe': '扣非PE', 'deductPeg': '扣非PEG',
                'yoyIncomeTotalGrowthRate': '营收累计同比',
                'yoyIncomeQuarterGrowthRate': '营收单季同比',
                'qoqIncomeGrowthRate': '营收季度环比', 'deductPegIncome': '营收PEG',
                'percentChangeToCalcDate': '首次公告至今涨幅',
                'percentChangeAfterWeek': '首次公告一周后涨幅',
                'percentChangeBeforeMonth': '首次公告一月前涨幅'
            }
            if 'groups' in df_outputs.columns:
                renamed_columns.__setitem__('groups', '分组')
            df_outputs = df_outputs.rename(columns=renamed_columns)

            df_outputs.to_excel(
                out_put_file,
                index=False,
                encoding='utf-8',
                columns=[renamed_columns[c] for c in columns if c in renamed_columns]
            )
            print('saved double growth to {}'.format(out_put_file))

        if to_db:
            bulk = []
            for tup in df_outputs.itertuples(index=False):
                one = UpdateOne(
                    {'code': tup.code, 'report_date': tup.report_date},
                    {'$set': tup._asdict()},
                    upsert=True
                )
                bulk.append(one)
            try:
                coll = settings.database.double_growth
                coll.bulk_write(bulk)
            except Exception as e:
                print(e)
            print('saved {} items to double_growth'.format(len(bulk)))
